import os
import cv2
import pickle
import numpy as np
from sklearn.cluster import DBSCAN
import insightface

class PhotoAlbumOrganizer:
    def __init__(self, album_path="input_images", feature_file="features.pkl"):
        self.album_path = album_path
        self.feature_file = feature_file
        self.model = insightface.app.FaceAnalysis()
        self.model.prepare(ctx_id=0, det_thresh=0.5)
        self.face_data = []
        
    def extract_features(self):
        """扫描相册目录并提取所有人脸特征"""
        if os.path.exists(self.feature_file):
            with open(self.feature_file, 'rb') as f:
                self.face_data = pickle.load(f)
            return
        
        for root, _, files in os.walk(self.album_path):
            for file in files:
                if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                    image_path = os.path.join(root, file)
                    print(f"Image: {image_path}")
                    image = cv2.imread(image_path)
                    if image is None:
                        continue
                        
                    faces = self.model.get(image)
                    for face in faces:
                        self.face_data.append({
                            "path": image_path,
                            "embedding": face.embedding,
                            "bbox": face.bbox.tolist()
                        })
                        print(face.bbox.tolist())
        
        with open(self.feature_file, 'wb') as f:
            pickle.dump(self.face_data, f)
    
    def cluster_faces(self, eps=0.5, min_samples=3):
        """对提取的特征进行DBSCAN聚类"""
        embeddings = np.array([data["embedding"] for data in self.face_data])
        clustering = DBSCAN(eps=eps, min_samples=min_samples).fit(embeddings)
        labels = clustering.labels_
        
        # 组织聚类结果
        clusters = {}
        for i, label in enumerate(labels):
            if label not in clusters:
                clusters[label] = []
            clusters[label].append(self.face_data[i])
        
        return clusters
    
    def visualize_results(self, clusters):
        """可视化聚类结果"""
        for label, faces in clusters.items():
            print(f"\nPerson {label} has {len(faces)} photos:")
            for face in faces[:3]:  # 每个类别显示前3张照片路径
                print(f"  - {face['path']}")

if __name__ == "__main__":
    organizer = PhotoAlbumOrganizer()
    organizer.extract_features()
    clusters = organizer.cluster_faces(eps=0.5, min_samples=3)
    organizer.visualize_results(clusters)
